import win32com.client
import time
import os
import sys
import openpyxl
import zipfile
from selenium import webdriver
# from selenium.webdriver.chrome.options import Options
from selenium.common.exceptions import NoSuchElementException

high_risk_port = ['135', '137', '138', '139', '445', '3389']


class DirScanner:
    def __init__(self, dir=None) -> None:
        self.dir = dir

    @classmethod
    def reportFileScan(cls, directory: str) -> list:
        '''找到所有report.html文件'''
        report_list = []
        for root, dirs, files in os.walk(directory):
            for file in files:
                # 获取文件路径
                if file == 'Report.html':
                    report_list.append(os.path.join(root, file))
        return report_list

    @classmethod
    def reportZipFileScan(cls, directory: str) -> list:
        '''找到所有report压缩包并返回'''
        report_zip_list = []
        for root, dirs, files in os.walk(directory):
            for file in files:
                # 获取文件路径
                if file.startswith('report_') and file.endswith('_html.zip'):
                    report_zip_list.append(os.path.join(root, file))
        return report_zip_list

    @classmethod
    def xlsxFileScan(cls, directory: str) -> list:
        '''查找目录下的所有excel表格'''
        xlsx_file_list = []
        for root, dirs, files in os.walk(directory):
            for file in files:
                # 获取文件路径
                if file.endswith('.xlsx'):
                    xlsx_file_list.append(os.path.join(root, file))
        return xlsx_file_list

    @classmethod
    def getSimilarXlsxFile(cls, *xlsx_files) -> list:

        basename_list = []
        same_unit_list = []
        if len(xlsx_files) == 0:
            return []
        directory = os.path.dirname(xlsx_files[0])
        for file in xlsx_files:
            basename_list.append(os.path.basename(file))
        
        for basename in basename_list:
            same_unit = []
            if not basename:
                continue
            for i in range(
                    basename_list.index(basename)+1,
                    len(basename_list)):
                if basename[:4] != basename_list[i][:4]:
                    continue
                same_unit.append(os.path.join(directory, basename_list[i]))
                basename_list[i] = None
            if len(same_unit) > 0:
                same_unit.insert(0, os.path.join(directory, basename))
                same_unit_list.append(same_unit)
        return same_unit_list


class Compresser:
    @classmethod
    def unzip(cls, *args):
        for file in args:
            with zipfile.ZipFile(file) as myzipfile:
                directory = file.split('.')[0]
                if not os.path.exists(directory):
                    print("%s正在解压" % file)
                    myzipfile.extractall(directory)
                    print("%s解压完毕" % file, end="\n\n")
                else:
                    print("路径%s已存在！不执行解压。。。" % directory, end="\n\n")


class Crawl:

    def __init__(
            self, url="") -> None:
        self.url = url
        chrome_options = webdriver.ChromeOptions()
        chrome_options.add_argument('--headless')  # 隐藏窗口
        chrome_options.add_argument('lang=zh_CN.UTF-8')
        UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) \
            AppleWebKit/537.36 (KHTML, like Gecko) Chrome/93.0.4577.63 Safari/537.36"
        chrome_options.add_experimental_option(
            "excludeSwitches", ['enable-automation', 'enable-logging'])  # 防止打印一些无用日志
        chrome_options.add_argument("User-Agent=" + UserAgent)
        self.driver = webdriver.Chrome(options=chrome_options)  # Chrome实例
        self.driver.get(self.url)  # driver.get方法将定位在给定的URL的网页
        time.sleep(0.5)

    def crawl(self, xlsx_file='test.xlsx'):
        driver = self.driver
        xlsx = openpyxl.load_workbook(xlsx_file)
        sheet = xlsx["Sheet"]
        text = ''
        for row in range(3, 1000):
            '''遍历漏扫报告中扫描过的ip地址条目'''
            driver.switch_to.frame(driver.find_element_by_id("mainFrame"))
            try:
                text = driver.find_element_by_xpath(
                    '//*[@id="table_13_content"]/tbody/tr[' + str(row) + ']').text
            except NoSuchElementException:
                break
            text = text.split()

            if len(text) < 8:
                '''判断扫描统计列表是否为空'''
                print('%s扫描统计列表为空！' % self.url, end="\n\n")
                break

            for col in range(1, 8):
                '''按行逐列写入漏洞数'''
                if col == 2:
                    print(text[col-1])  # 打印当前正在扫描的ip地址
                sheet.cell(row=row-1, column=col).value = \
                    int(text[col-1]) if text[col-1].isdigit() \
                    else text[col-1]  # 按列赋值

            driver.find_element_by_xpath(
                '//*[@id="table_13_content"]/tbody/tr[' + str(row) + ']/td[2]/a').click()
            time.sleep(0.2)
            handle = driver.window_handles
            try:
                driver.switch_to.window(handle[1])
            except:
                sys.stderr.write("设备（ip）详情页还没打开，0.5秒后再次尝试。")
                time.sleep(0.5)
                driver.switch_to.window(handle[1])
                
            driver.switch_to.frame(driver.find_element_by_id("mainFrame"))
            port_list = ''
            for t in range(4, 100):
                port = '0'
                try:
                    port_and_bug = driver.find_element_by_xpath(
                        '//*[@id="table_' + str(t) + '_content"]/tbody/tr[2]/th/div').text
                    port = port_and_bug.split('__')[0]
                except NoSuchElementException:
                    try:
                        port_and_bug = driver.find_element_by_xpath(
                            '//*[@id="table_' + str(t) + '_content"]/tbody/tr[2]/td/div').text
                        port = port_and_bug.split('__')[0]
                    except NoSuchElementException:
                        break

                if port != '0' and port in high_risk_port:
                    port_list = port_list + '、' + port

            for t in range(2, 100):
                port = '0'
                try:
                    port_and_bug = driver.find_element_by_xpath(
                        '//*[@id="table_5_content"]/tbody/tr[' + str(t) + ']/td/div').text
                    port = port_and_bug.split('__')[0]
                except NoSuchElementException:
                    break

                if port != '0' and port in high_risk_port:
                    port_list = port_list + '、' + port

            port_list = port_list[1:]
            sheet.cell(row=row-1, column=8).value = port_list
            driver.close()
            driver.switch_to.window(handle[0])
        driver.quit()  # 关闭浏览器

        # Xlsx.sumByColumn(sheet, 3, 8)
        xlsx.save(xlsx_file)
        xlsx.close()


class Common:
    @classmethod
    def reportDirInput(cls):
        FLAG = False
        report_zip_source_dir = "."
        TIMES = 3
        while not FLAG:
            if TIMES <= 0:
                print("目录错误输入次数过多！程序退出。。。")
                sys.exit(1)
            report_zip_source_dir = input("请输入漏扫报告所在的路径：")
            TIMES -= 1
            if os.path.isdir(report_zip_source_dir):
                return report_zip_source_dir
            print("路径不存在！！", end="")

    @classmethod
    def xlsxDirInput(cls):
        to_dir = "."
        TIMES = 3
        while 1:
            if TIMES <= 0:
                print("目录错误输入次数过多！程序退出。。。")
                sys.exit(1)
            to_dir = input("请输入表格存入的目标路径：")
            TIMES -= 1
            if os.path.isdir(to_dir):
                return to_dir
            print("路径不存在！！", end="")


class Xlsx:
    def __init__(self, filename: str) -> None:
        self.wb = openpyxl.load_workbook(filename)
    
    
    @classmethod
    def sumByColumn(cls, sheet, col_from=3, col_to=7) -> int:
        '''
        将sheet中的数据在最后一行后进行汇总
        '''
        lastline = sheet.max_row
        next_lastline = lastline + 1
        for col in range(col_from, col_to+1):
            '''汇总统计不同级别漏洞总数'''
            sum_leaks = 0
            for row in range(2, next_lastline):
                num = sheet.cell(row=row, column=col).value
                if num is not None:
                    try:
                        sum_leaks += int(num)
                    except TypeError:
                        print('TypeError num: %s', num)
                        break
            sheet.cell(row=next_lastline, column=col).value = sum_leaks
        # 将最后一行行号写入汇总行的第二列
        sheet.cell(row=next_lastline, column=2).value = \
            sheet.cell(row=lastline, column=1).value
        return next_lastline

    @classmethod
    def initReportWorksheet(cls, filename: str) -> str:
        '''初始化每一个漏扫报告对应的表格'''
        xlsx = openpyxl.Workbook()

        sheet = xlsx["Sheet"]
        sheet.cell(row=1, column=1).value = '序号'

        sheet.cell(row=1, column=2).value = '主机IP'
        sheet.column_dimensions['B'].width = 20.0

        sheet.cell(row=1, column=3).value = '漏洞总数'

        sheet.cell(row=1, column=4).value = '高危漏洞数'
        sheet.column_dimensions['D'].width = 10.0

        sheet.cell(row=1, column=5).value = '中危漏洞数'
        sheet.column_dimensions['E'].width = 10.0

        sheet.cell(row=1, column=6).value = '低危漏洞数'
        sheet.cell(row=1, column=7).value = '信息漏洞'
        sheet.cell(row=1, column=8).value = '高危端口'

        xlsx.save(filename)
        xlsx.close()
        return filename
    
    @classmethod
    def initSummaryWorksheet(cls, sheet: openpyxl.Workbook.worksheets):
        sheet.cell(row=1, column=1).value = '单位名'
        sheet.column_dimensions['A'].width = 30.0

        sheet.cell(row=1, column=2).value = '在线主机数量'
        sheet.column_dimensions['B'].width = 15.0

        sheet.cell(row=1, column=3).value = '漏洞总数'
        sheet.column_dimensions['C'].width = 15.0

        sheet.cell(row=1, column=4).value = '高危漏洞数'
        sheet.column_dimensions['D'].width = 15.0

        sheet.cell(row=1, column=5).value = '中危漏洞数'
        sheet.column_dimensions['E'].width = 15.0

        sheet.cell(row=1, column=6).value = '低危漏洞数'
        sheet.column_dimensions['F'].width = 15.0
        
        sheet.cell(row=1, column=7).value = '信息漏洞数'
        sheet.column_dimensions['G'].width = 15.0

    @classmethod
    def getLastLineNum(cls, filename: str) -> int:
        '''获取excel表的最后一行'''
        xlsx_file = openpyxl.load_workbook(filename)
        sheet = xlsx_file["Sheet"]
        return sheet.max_row
    
    @classmethod
    def mergeSameNameFile(cls, *same_unit_list):
        '''融合同站段的不同表'''
        for same_unit in same_unit_list:
            if len(same_unit) == 0:
                continue
            # 以同名列表中的第一作为同名汇总表
            to_xlsx = openpyxl.load_workbook(same_unit[0])
            to_xlsx_sheet = to_xlsx["Sheet"]
            to_sheet_startline = to_xlsx_sheet.max_row + 1
            for source_xlsx in same_unit[1:]:
                # 打开源excel表
                s_xlsx = openpyxl.load_workbook(source_xlsx)  
                source_xlsx_sheet = s_xlsx["Sheet"]
                source_sheet_lastline = source_xlsx_sheet.max_row
                for row in range(2, source_sheet_lastline+1):
                    to_xlsx_sheet.cell(row=to_sheet_startline, column=1).value = to_sheet_startline - 1
                    for col in range(2,9):
                        to_xlsx_sheet.cell(row=to_sheet_startline, column=col).value = \
                            source_xlsx_sheet.cell(row=row, column=col).value
                    to_sheet_startline += 1
                s_xlsx.close()
                os.remove(source_xlsx)
            to_xlsx.save(same_unit[0])
            to_xlsx.close()


    @classmethod
    def sheetAppend(cls, s_sheet, d_sheet, s_row=1, d_row=1):
        '''
        s_sheet: 源sheet
        d_sheet: 目标sheet
        s_row: 源起始行
        d_row: 目标起始行
        '''
        drow = d_row
        for srow in range(s_row, s_sheet.max_row+1):
            for col in range(1, s_sheet.max_column+1):
                d_sheet.cell(row=drow, column=col).value = \
                    s_sheet.cell(row=srow, column=col).value
            drow += 1

if __name__ == "__main__":
    # 标准输入漏扫报告压缩包所在的路径
    source_dir = Common.reportDirInput()
    # 标准输入xlsx存入的路径
    to_xlsx_dir = Common.xlsxDirInput()
    
    # 将所有站段漏扫数据汇总到同一个工作簿
    summary_workbook_file = os.path.join(to_xlsx_dir, "总表.xlsx")
    if os.path.exists(summary_workbook_file):
        sys.stderr.write("%s 已存在！" % summary_workbook_file)
        sys.exit()
        
    report_zip_list = \
        DirScanner.reportZipFileScan(source_dir)
    if len(report_zip_list) <= 0:
        print("输入路径不存在压缩包！程序退出。。。")
        sys.exit()

    Compresser.unzip(*report_zip_list)
    report_file_list = \
        DirScanner.reportFileScan(source_dir)
    if len(report_file_list) <= 0:
        print("输入路径不存Report.html文件！程序退出。。。")
        sys.exit()

    # 爬取每一个漏扫报告中的数据，使一个漏扫报告一个表格
    for report_url in report_file_list:
        print("正在扫描：%s 。。。" % report_url)
        report_name = report_url.split('\\')[-3]
        if len(report_name.split('_')) == 5:
            xlsx_filename = report_name.split('_')[1] + '.xlsx'
        else:
            xlsx_filename = '_'.join(report_name.split('_')[1:3])
        to_xlsx = os.path.join(to_xlsx_dir, xlsx_filename)
        xlsx = Xlsx.initReportWorksheet(to_xlsx)
        crawl = Crawl(url=report_url)
        crawl.crawl(xlsx_file=to_xlsx)
        print("%s 扫描完成，excel表格已生成！" % report_url, end="\n\n")

    # 汇总同站段扫描报告
    xlsx_file_list = DirScanner.xlsxFileScan(to_xlsx_dir)
    similarXlsxFile = DirScanner.getSimilarXlsxFile(*xlsx_file_list)
    Xlsx.mergeSameNameFile(*similarXlsxFile)
    
    summary_wb = openpyxl.Workbook()
    summary_wb.remove(summary_wb["Sheet"])
    xlsx_file_list = DirScanner.xlsxFileScan(to_xlsx_dir)
    for i in range(len(xlsx_file_list)):
        xlsx_file = xlsx_file_list[i]
        wb = ''
        try:
            wb = openpyxl.load_workbook(xlsx_file, read_only=True)
        except zipfile.BadZipFile:
            sys.stderr.write("%s 文件已损坏！" % xlsx_file)
        basename = xlsx_file.split("\\")[-1].split(".")[0]
        summary_wb.create_sheet(basename, i)
        try:
            Xlsx.sheetAppend(wb["Sheet"], summary_wb[basename])
        except KeyError:
            sys.stderr.write("%s 不存在名为\'Sheet\'的sheet" % xlsx_file)
            wb.close()
            continue
        wb.close()
    
    # 将工作簿中所有的表格汇总为一个表
    summary_wb.create_sheet("汇总", 0)
    summary_worksheet = summary_wb["汇总"]
    # 初始化汇总表格
    Xlsx.initSummaryWorksheet(summary_worksheet)
    sheets_num = len(summary_wb.worksheets)
    for index in range(1, sheets_num):
        # 汇总数据
        sheet = summary_wb.worksheets[index]
        lastline = Xlsx.sumByColumn(sheet)       
        summary_worksheet.cell(row=index+1, column=1).value = \
            summary_wb.worksheets[index].title
        for col in range(2, 8):
            summary_wb["汇总"].cell(row=index+1, column=col).value = \
                summary_wb.worksheets[index].cell(
                    row=lastline, column=col
                ).value
    
    summary_wb.save(summary_workbook_file)
    summary_wb.close()
    
    excel = win32com.client.Dispatch('Excel.Application')
    """
    0代表隐藏对象，但可以通过菜单再显示
    -1代表显示对象
    2代表隐藏对象，但不可以通过菜单显示，只能通过VBA修改为显示状态
    """
    excel.Visible = -1

    # 打开excel
    myBook = excel.Workbooks.Open(summary_workbook_file)